import uvicorn
from fastapi import FastAPI
from typing import Annotated
from fastapi import File, Form, UploadFile
from langchain_helper import *
from pydantic import BaseModel
import os
import json

## python -m uvicorn main:app --reload

file_upload_path = "./uploads"
app = FastAPI()


@app.get("/")
def read_root():
    return {"Hello simon"}


@app.post("/upload")
def upload_pdf(
        file: Annotated[UploadFile, File()],
        index_name: Annotated[str, Form()]
):
    file_upload_target_path = os.path.join(file_upload_path, file.filename)
    with open(file_upload_target_path, "wb") as f:
        contents = file.file.read()
        f.write(contents)
    load_pdf_and_save_to_index(file_upload_target_path, index_name)
    return {"filename": file.filename, "index_name": index_name}


class Query(BaseModel):
    index_name: str
    query: str


@app.post("/query")
def query_index(requset: Query):
    # json_str=json.dumps(data)
    # json_data=json.loads(json_str)
    # index_name = json_data['index_name']
    # query=json_data['query']
    index_name = requset.index_name
    query = requset.query
    index = load_index(index_name)
    ans = query_index_lc(index, query)
    return {"ans": ans, "index_name": index_name, "query": query}


if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=9001)
#   query = Query(index_name = "simon_index001",query = "郎亚坤的用工形式")
#   ans = query_index(query)
#   print(ans)
